import jpype
from jpype import java
from jpype import javax
from kazoo.client import KazooClient
import time
import json
import sqlite3
import datetime
import configparser
import xlsxwriter
import os
import logging.handlers

os.makedirs('./logs/', exist_ok=True)
os.makedirs('./warehouse', exist_ok=True)
loghandler = logging.handlers.TimedRotatingFileHandler(filename="logs/monitor.log", when="h", interval=24,
                                                       backupCount=30)
formatter = logging.Formatter(fmt='%(asctime)s,%(lineno)s,%(levelname)s:%(message)s', datefmt='%Y-%m-%d %H:%M:%S')
loghandler.setFormatter(formatter)
loghandler.suffix = "%Y-%m-%d_%H-%M-%S"
logging.basicConfig(level=logging.INFO)
log = logging.getLogger(__name__)
log.addHandler(loghandler)


class Main():
    def load_config(self):
        self.conf = configparser.ConfigParser()
        self.conf.read('./config.ini')
        self.byte_in = "kafka.server:type=BrokerTopicMetrics,name=BytesInPerSec,topic={}"
        self.byte_out = 'kafka.server:type=BrokerTopicMetrics,name=BytesOutPerSec,topic={}'
        self.attribute = "OneMinuteRate"
        self.count = 'Count'
        self.sql_mon_minute = 'insert into {}_mon_minute values (?,?,?,?,?,?);'
        self.sql_mon_minute_statistics = 'insert into {}_mon_minute_statistics values (?,?,?,?,?);'
        self.units = ['b', 'kb', 'mb', 'gb']

    def load_jmx_cli(self):
        # init jmx client
        jvmPath = jpype.getDefaultJVMPath()
        jpype.startJVM(jvmPath)
        jhash = java.util.HashMap()
        for i in range(len(self.jmxurls)):
            url = self.jmxurls[i]['jmxrmi']
            jmxurl = javax.management.remote.JMXServiceURL(url)
            jmxsoc = javax.management.remote.JMXConnectorFactory.connect(jmxurl, jhash)
            connection = jmxsoc.getMBeanServerConnection()
            self.jmxurls[i]['client'] = connection

    def load_url(self):
        kafka_path = self.conf.get('zookeeper', 'kafka_path')
        zk_host = self.conf.get('zookeeper', 'host')
        zk = KazooClient(hosts=zk_host, read_only=True)
        zk.start()
        self.jmxurls = []
        id_path = '{}/brokers/ids'.format(kafka_path)
        jmxrmi = 'service:jmx:rmi:///jndi/rmi://{}:{}/jmxrmi'
        # load jmxrmi dict
        for id in zk.get_children(id_path):
            data = json.loads(zk.get('{}/{}'.format(id_path, id))[0].decode('utf-8'))
            info = {'id': int(id),
                    'jmxrmi': jmxrmi.format(data.get('host'), data.get('jmx_port')),
                    'host': data.get('host')}
            self.jmxurls.append(info)
        # load all topic and partitions
        self.topics = []
        topic_path = '{}/brokers/topics'.format(kafka_path)
        topics = self.conf.get('common', 'topics').split(',')
        for topic in zk.get_children(topic_path):
            if topic.startswith('__') or topic not in topics:
                continue
            partitions = zk.get_children('{}/{}/partitions'.format(topic_path, topic))
            info = {'name': topic, 'partition_size': len(partitions), 'broker': set()}
            for partition in partitions:
                partition_path = '{}/{}/partitions/{}/state'.format(topic_path, topic, partition)
                state = json.loads(zk.get(partition_path)[0].decode('utf-8'))
                info['broker'].add(state['leader'])
            self.topics.append(info)
        zk.stop()
        zk.close()

    def load_sqlite(self):
        self.cli = sqlite3.connect('monitor.db')
        cursor = self.cli.cursor()
        cursor.execute('select tbl_name from sqlite_master;')
        minute_sql = 'create table {}_mon_minute (topic text, host text, stime text, inflow text, outflow text, msg_count text);'
        minute_statistics_sql = 'create table {}_mon_minute_statistics (topic text, stime text, inflow text, outflow text, msg_increase text);'
        topics = self.conf.get('common', 'topics').split(',')
        # 在 topic 里面 但不在 table 里面的
        tables = set([_[0].split('_')[0] for _ in cursor.fetchall()])
        create = set(topics).difference(tables)
        for table in create:
            cursor.execute(minute_sql.format(table))
            cursor.execute(minute_statistics_sql.format(table))
        cursor.close()

    def human(self, num, unit=0):
        if num > 1024:
            num = num / 1024
            unit += 1
            return self.human(num, unit)
        else:
            return round(num, 2), unit

    def handle(self):
        now = datetime.datetime.now().strftime('%Y%m%d%H%M')
        mon_minutes = []
        mon_minute_statistics = []
        cursor = self.cli.cursor()
        for topic in self.topics:
            in_rate_all = 0
            out_rate_all = 0
            msg_count_all = 0
            for broker in topic['broker']:
                for cli in self.jmxurls:
                    if broker == cli['id']:
                        try:
                            in_rate = cli['client'].getAttribute(
                                javax.management.ObjectName(self.byte_in.format(topic['name'])),
                                self.attribute)
                            msg_count = cli['client'].getAttribute(
                                javax.management.ObjectName(self.byte_in.format(topic['name'])), self.count)
                            out_rate = cli['client'].getAttribute(
                                javax.management.ObjectName(self.byte_out.format(topic['name'])), self.attribute)
                            mon_minutes.append([topic['name'], cli['host'], now, in_rate, out_rate, msg_count])
                            in_rate_all = in_rate_all + in_rate
                            out_rate_all = out_rate_all + out_rate
                            msg_count_all = msg_count_all + msg_count
                            # topic,broker,in_flow,out_flow,msg_count,
                            log.info('%s,%s,%s,%s,%s', topic['name'], broker, in_rate, out_rate, msg_count)
                        except Exception as err:
                            log.error('%s:%s:not find', topic['name'], broker)
            in_rate_all_h, in_rate_all_unit = self.human(in_rate_all)
            out_rate_all_h, out_rate_all_unit = self.human(out_rate_all)
            in_all = '{}-{}'.format(in_rate_all_h, self.units[in_rate_all_unit])
            out_all = '{}-{}'.format(out_rate_all_h, self.units[out_rate_all_unit])
            mon_minute_statistics.append([topic['name'], now, in_all, out_all, msg_count_all])
            cursor.executemany(self.sql_mon_minute.format(topic['name']), mon_minutes)
            cursor.executemany(self.sql_mon_minute_statistics.format(topic['name']), mon_minute_statistics)
        self.cli.commit()

    def show(self):
        cursor = self.cli.cursor()
        file_name = datetime.datetime.now().strftime('%Y%m%d%H%M')
        file_name = "{}/{}.xlsx".format('./warehouse', file_name)
        workbook = xlsxwriter.Workbook(file_name)
        for topic in self.topics:
            package_data = self.show_package_data(topic['name'], cursor)
            print(package_data)
            all_in_flow = package_data['all_in_flow']
            use_times = package_data['use_times']
            in_flow_unit = package_data['in_flow_unit']
            worksheet = workbook.add_worksheet(topic['name'])
            bold = workbook.add_format({'bold': 0.1})
            use_times.insert(0, 'time')
            worksheet.write_row('A1', use_times, bold)
            count = 2
            chart_col = workbook.add_chart({'type': 'line'})
            # 将key 全部取出后 排序
            for t in sorted(list(all_in_flow.keys())):
                v = all_in_flow.get(t)
                v.insert(0, '{}:{}'.format(t[:2], t[2:]))
                worksheet.write_row('A{}'.format(count), v, bold)
                count += 1
            colors = ['#FF3333', '#00FF66', '#FFFF33', '#66FFFF', '#000000', '#0033FF']
            for i in range(len(use_times) - 1):
                series = {
                    'name': [topic['name'], 0, i + 1],
                    'categories': [topic['name'], 1, 0, len(all_in_flow.keys()) + 1, 0],
                    'values': [topic['name'], 1, i + 1, len(all_in_flow.keys()) + 1, i + 1],
                    'line': {'color': colors[i], 'width': 1.25},
                }
                chart_col.add_series(series)
            # # 设置图表的title 和 x，y轴信息
            chart_col.set_title({'name': '{} 五天内流入量'.format(topic['name'])})
            chart_col.set_x_axis({'name': '时间'})
            chart_col.set_y_axis({'name': '流入量({}/second)'.format(in_flow_unit)})
            # # 设置图表的风格
            chart_col.set_style(20)
            # # 把图表插入到worksheet并设置偏移
            worksheet.insert_chart('A1', chart_col, {'x_scale': 3, 'y_scale': 2})
        workbook.close()
        os.system('expect scp.exp {} /sc_bin/check_kafka_log/kafka_log_mon '.format(file_name))
        os.remove(file_name)

    def package_time(self):
        start = datetime.datetime.strptime('0000', '%H%M')
        use_times = {(start + datetime.timedelta(minutes=_)).strftime('%H%M'): [] for _ in range(1440)}
        return use_times

    def show_package_data(self, topic, cursor):
        # TODO 这里假设的是所有数据都是存在的 但是如果某一天的数据断了，是无法显示出来的只会由有一天的数据补上，这就造成了数据的错位
        now = datetime.datetime.now()
        # 组装一天每分钟的key 在后期添加键
        all_in_flow = self.package_time()
        use_times = []
        template_sql = "select stime, inflow from {}_mon_minute_statistics where topic='{}' and stime like '{}%' order by stime;"
        in_unit_set = set()
        for i in range(0, 6):
            use_time = now + datetime.timedelta(days=-i)
            use_time = use_time.strftime('%Y%m%d')
            use_times.append(use_time)
            sql = template_sql.format(topic, topic, use_time)
            cursor.execute(sql)
            datas = {}
            for data in cursor.fetchall():
                in_flow_unit = data[1].split('-')
                in_flow = round(float(in_flow_unit[0]), 2)
                datas[data[0][-4:]] = in_flow
                in_unit_set.add(in_flow_unit[1])

            for k, v in all_in_flow.items():
                v.append(datas.get(k, 0))
                all_in_flow[k] = v
        result = {
            'all_in_flow': all_in_flow,
            'use_times': use_times,
            'in_flow_unit': in_unit_set.pop() if len(in_unit_set) == 1 else 'error'
        }
        return result

    def start(self):
        self.load_config()
        self.load_sqlite()
        self.load_url()
        self.load_jmx_cli()
        isshow = datetime.datetime.strptime('1100', '%H%M')
        while True:
            self.handle()
            now = datetime.datetime.now()
            if now.hour == isshow.hour and now.minute == isshow.minute:
                self.show()
            time.sleep(60)
        # jpype.shutdownJVM()
        # self.cli.close()


if __name__ == '__main__':
    main = Main()
    main.start()
